//
// Created by DELL on 2022/4/22.
//
#pragma once

#include "../Utils.h"
#include "../VulkanDefine.h"

BEGIN_NS_VULKAN

enum UniformType
{
    kU_void,
    kU_float,
    KU_vec2,
    KU_vec3,
    KU_vec4,
    kU_double,
    kU_mat3,
    kU_mat4,
    kU_int8,
    kU_int16,
    kU_int32,
    kU_uint8,
    kU_uint16,
    kU_uint32,
    kU_bool,
    kU_atomic,
    kU_sampler,
    kU_struct,
};
struct UniformField
{
    UniformType type;
    int offset;
};
static constexpr int UniformFieldTypeSize[] {
    0,sizeof(float), sizeof(math::Vec2), sizeof(math::Vec3), sizeof(math::Vec4), sizeof(double),
    sizeof(math::Mat3), sizeof(math::Mat4), sizeof(int8_t), sizeof(int16_t), sizeof(int32_t),
    sizeof(uint8_t), sizeof(uint16_t), sizeof(uint32_t), sizeof(bool), 0, 0, 0, 0 };

using UniformFieldSet = std::unordered_map<std::string, UniformField>;

struct UniformBlock
{
    std::string name;
    int binding{0};
    int size{0};
    VkShaderStageFlags stages{0};
    UniformFieldSet fields;
};

struct Variable
{
    std::string name;
    UniformType type{kU_sampler};
    int binging{0};
    VkShaderStageFlags stages{0};
};

struct DescriptorSetLayout
{
    std::vector<UniformBlock> blocks;
    std::unordered_map<std::string, Variable> samples;

    VkDescriptorSetLayout createVkLayout(VkDevice device);
};

class DescriptorSet
{
public:
    ~DescriptorSet();
    DescriptorSet(const DescriptorSet&) = delete;
    DescriptorSet& operator=(const DescriptorSet&) = delete;

    VkDescriptorSet& getDescriptorSet() { return this->descriptorSet; }

    static DescriptorSetRef create(DeviceRef& device, VkDescriptorSetLayout layout, VkDescriptorPool pool);

    void addBuffer(UniformBufferRef& uniform, int binding);
    void addSampler(TextureRef& texture, VkSampler sampler, int binding);
    void addSampler(VkImageView imageView, VkSampler sampler, int binding);
    void build();

private:
    explicit DescriptorSet(DeviceRef device);

    DeviceRef device;
    VkDescriptorSet descriptorSet{};

    std::vector<VkWriteDescriptorSet> writeDescriptorSet;
    std::vector<VkDescriptorBufferInfo> bufferInfos;
    std::vector<VkDescriptorImageInfo> imageInfos;
};

END_NS_VULKAN